#!/usr/bin/env python3

import socket
import sys
from time import sleep

from scapy.all import (
    IP,
    UDP,
    Ether,
    FieldLenField,
    IntField,
    IPOption,
    Packet,
    PacketListField,
    ShortField,
    get_if_hwaddr,
    get_if_list,
    sendp
)
from scapy.layers.inet import _IPOption_HDR


def get_if():
    ifs = get_if_list()
    iface = None  # 默认接口
    for i in ifs:
        if "eth0" in i:
            iface = i
            break
    if not iface:
        print("Cannot find eth0 interface")
        exit(1)
    return iface

class SwitchTrace(Packet):
    fields_desc = [ IntField("swid", 0),
                    IntField("qdepth", 0)]
    def extract_padding(self, p):
        return "", p

class IPOption_MRI(IPOption):
    name = "MRI"
    option = 31
    fields_desc = [_IPOption_HDR,
                   FieldLenField("length", None, fmt="B", length_of="swtraces", adjust=lambda pkt, l: l * 2 + 4),
                   ShortField("count", 0),
                   PacketListField("swtraces", [], SwitchTrace, count_from=lambda pkt: (pkt.count * 1))]


def main():
    if len(sys.argv) != 5:
        print('Usage: send.py <source_ip> <destination> "<message>" <count>')
        exit(1)

    src_ip = sys.argv[1]
    dst_ip = socket.gethostbyname(sys.argv[2])
    message = sys.argv[3]
    count = int(sys.argv[4])
    iface = get_if()
    pkt = Ether(src=get_if_hwaddr(iface), dst="ff:ff:ff:ff:ff:ff") / IP(src=src_ip, dst=dst_ip, options=IPOption_MRI(count=0, swtraces=[])) / UDP(dport=4321, sport=1234) / message
    pkt.show2()

    try:
        for i in range(count):
            sendp(pkt, iface=iface)
            sleep(1)
    except KeyboardInterrupt:
        raise

if __name__ == '__main__':
    main()
